// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

// Array accessor types for primitive/C-type-based arrays, such as numbers,
// boolean, and temporal types.

#pragma once

#include <cstdint>
#include <memory>

#include <nebula/array/array_base.h>
#include <nebula/array/data.h>
#include <nebula/core/stl_iterator.h>
#include <nebula/types/type.h>
#include <nebula/types/type_fwd.h>  // IWYU pragma: export
#include <nebula/types/type_traits.h>
#include <nebula/bits/bit_util.h>
#include <turbo/base/macros.h>

namespace nebula {

    /// Concrete Array class for boolean data
    class TURBO_EXPORT BooleanArray : public PrimitiveArray {
    public:
        using TypeClass = BooleanType;
        using IteratorType = stl::ArrayIterator<BooleanArray>;

        explicit BooleanArray(const std::shared_ptr<ArrayData> &data);

        BooleanArray(int64_t length, const std::shared_ptr<Buffer> &data,
                     const std::shared_ptr<Buffer> &null_bitmap = nullptr,
                     int64_t null_count = kUnknownNullCount, int64_t offset = 0);

        bool get_value(int64_t i) const {
            return bit_util::get_bit(reinterpret_cast<const uint8_t *>(raw_values_),
                                    i + data_->offset);
        }


        bool value(int64_t i) const {
            return get_value(i);
        }

        bool get_view(int64_t i) const { return value(i); }

        std::optional<bool> operator[](int64_t i) const { return *IteratorType(*this, i); }

        /// \brief Return the number of false (0) values among the valid
        /// values. turbo::Result is not cached.
        int64_t false_count() const;

        /// \brief Return the number of true (1) values among the valid
        /// values. turbo::Result is not cached.
        int64_t true_count() const;

        IteratorType begin() const { return IteratorType(*this); }

        IteratorType end() const { return IteratorType(*this, length()); }

    protected:
        using PrimitiveArray::PrimitiveArray;
    };

    /// \addtogroup numeric-arrays
    ///
    /// @{

    /// \brief Concrete Array class for numeric data with a corresponding C type
    ///
    /// This class is templated on the corresponding DataType subclass for the
    /// given data, for example NumericArray<Int8Type> or NumericArray<Date32Type>.
    ///
    /// Note that convenience aliases are available for all accepted types
    /// (for example Int8Array for NumericArray<Int8Type>).
    template<typename TYPE>
    class NumericArray : public PrimitiveArray {
    public:
        using TypeClass = TYPE;
        using value_type = typename TypeClass::c_type;
        using IteratorType = stl::ArrayIterator<NumericArray<TYPE>>;

        explicit NumericArray(const std::shared_ptr<ArrayData> &data) : PrimitiveArray(data) {}

        // Only enable this constructor without a type argument for types without additional
        // metadata
        template<typename T1 = TYPE>
        NumericArray(enable_if_parameter_free<T1, int64_t> length,
                     const std::shared_ptr<Buffer> &data,
                     const std::shared_ptr<Buffer> &null_bitmap = nullptr,
                     int64_t null_count = kUnknownNullCount, int64_t offset = 0)
                : PrimitiveArray(TypeTraits<T1>::type_singleton(), length, data, null_bitmap,
                                 null_count, offset) {}

        const value_type *raw_values() const {
            return reinterpret_cast<const value_type *>(raw_values_) + data_->offset;
        }

        value_type get_value(int64_t i) const { return raw_values()[i]; }

        value_type value(int64_t i) const { return raw_values()[i]; }

        // For API compatibility with BinaryArray etc.
        value_type get_view(int64_t i) const { return value(i); }

        std::optional<value_type> operator[](int64_t i) const {
            return *IteratorType(*this, i);
        }

        IteratorType begin() const { return IteratorType(*this); }

        IteratorType end() const { return IteratorType(*this, length()); }

    protected:
        using PrimitiveArray::PrimitiveArray;
    };

    /// DayTimeArray
    /// ---------------------
    /// \brief Array of Day and Millisecond values.
    class TURBO_EXPORT DayTimeIntervalArray : public PrimitiveArray {
    public:
        using TypeClass = DayTimeIntervalType;
        using IteratorType = stl::ArrayIterator<DayTimeIntervalArray>;

        explicit DayTimeIntervalArray(const std::shared_ptr<ArrayData> &data);

        DayTimeIntervalArray(const std::shared_ptr<DataType> &type, int64_t length,
                             const std::shared_ptr<Buffer> &data,
                             const std::shared_ptr<Buffer> &null_bitmap = nullptr,
                             int64_t null_count = kUnknownNullCount, int64_t offset = 0);

        DayTimeIntervalArray(int64_t length, const std::shared_ptr<Buffer> &data,
                             const std::shared_ptr<Buffer> &null_bitmap = nullptr,
                             int64_t null_count = kUnknownNullCount, int64_t offset = 0);

        TypeClass::DayMilliseconds get_value(int64_t i) const;

        TypeClass::DayMilliseconds value(int64_t i) const { return get_value(i); }

        // For compatibility with Take kernel.
        TypeClass::DayMilliseconds get_view(int64_t i) const { return get_value(i); }

        IteratorType begin() const { return IteratorType(*this); }

        IteratorType end() const { return IteratorType(*this, length()); }

        std::optional<TypeClass::DayMilliseconds> operator[](int64_t i) const {
            return *IteratorType(*this, i);
        }

        int32_t byte_width() const { return sizeof(TypeClass::DayMilliseconds); }

        const uint8_t *raw_values() const { return raw_values_ + data_->offset * byte_width(); }
    };

    /// \brief Array of Month, Day and nanosecond values.
    class TURBO_EXPORT MonthDayNanoIntervalArray : public PrimitiveArray {
    public:
        using TypeClass = MonthDayNanoIntervalType;
        using IteratorType = stl::ArrayIterator<MonthDayNanoIntervalArray>;

        explicit MonthDayNanoIntervalArray(const std::shared_ptr<ArrayData> &data);

        MonthDayNanoIntervalArray(const std::shared_ptr<DataType> &type, int64_t length,
                                  const std::shared_ptr<Buffer> &data,
                                  const std::shared_ptr<Buffer> &null_bitmap = nullptr,
                                  int64_t null_count = kUnknownNullCount, int64_t offset = 0);

        MonthDayNanoIntervalArray(int64_t length, const std::shared_ptr<Buffer> &data,
                                  const std::shared_ptr<Buffer> &null_bitmap = nullptr,
                                  int64_t null_count = kUnknownNullCount, int64_t offset = 0);

        TypeClass::MonthDayNanos get_value(int64_t i) const;

        TypeClass::MonthDayNanos value(int64_t i) const { return get_value(i); }

        // For compatibility with Take kernel.
        TypeClass::MonthDayNanos get_view(int64_t i) const { return get_value(i); }

        IteratorType begin() const { return IteratorType(*this); }

        IteratorType end() const { return IteratorType(*this, length()); }

        std::optional<TypeClass::MonthDayNanos> operator[](int64_t i) const {
            return *IteratorType(*this, i);
        }

        int32_t byte_width() const { return sizeof(TypeClass::MonthDayNanos); }

        const uint8_t *raw_values() const { return raw_values_ + data_->offset * byte_width(); }
    };

    /// @}

}  // namespace nebula
